import sys
import os
# 添加 project/ 目录到 sys.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import cv2
import numpy as np
import argparse
from tictoc import TicToc
import multiprocessing
from tqdm import tqdm

extra_topics = ["SceneColor", "SceneDepth", "GBufferA", "GBufferB", "GBufferC", "GBufferD", "Velocity",
                "GBufferSSAO", "CustomStencil"]


def main(args):
    # 统计耗时
    cost = TicToc("gbuffers图片处理")
    assert os.path.exists(args.data_path)
    frames = os.listdir(args.data_path)
    frames.sort(key=lambda x: x)
    files = []
    for dir in frames:
        if dir[:2] == '__':  # '__'
            dir = os.path.join(args.data_path, dir)
            if os.path.isdir(dir):
                files.append(dir)

    process_size = len(files)
    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


def main_worker(frames_path, args):
    camera_topics = ["camera75", "camera77", "camera80", "camera81"]

    output_paths = []
    output_path = os.path.join(frames_path, 'camera_epe')
    output_paths.append(output_path)
    for topic in camera_topics:
        output_path = os.path.join(frames_path, 'camera_epe', topic)
        output_paths.append(output_path)
        output_path = os.path.join(frames_path, 'camera_epe', topic, 'GBuffers')
        output_paths.append(output_path)
    for output_path in output_paths:
        if output_path is not None and not os.path.exists(output_path):
            print("新建目录<{}>".format(output_path))
            os.makedirs(output_path, exist_ok=True)

    for topic in camera_topics:
        path = os.path.join(frames_path, 'camera_extra', topic, extra_topics[0])
        file_generator = (os.path.join(root, file) for root, dirs, files in os.walk(os.path.abspath(path)) for file in
                          files)
        # 使用 tqdm 包装生成器，并设置 postfix 参数
        lines = path.strip().split('/')
        for file_path in tqdm(file_generator, postfix=os.path.join(lines[-4], lines[-2]) + "/ gbuffers图片处理中..."):
            # 确保文件是图片文件，这里假设你的图片文件是 .jpg 或 .png 格式
            if file_path.lower().endswith(('.jpg', '.png')):
                name = os.path.basename(file_path).split(".")[0]
                gt_labels = material_from_gbuffers(os.path.join(frames_path, 'camera_extra', topic), name)
                save_path = os.path.join(frames_path, 'camera_epe', topic, 'GBuffers', name + ".npz")
                np.savez_compressed(save_path, gt_labels)


def material_from_gbuffers(file_path, name):
    gbuff = {}
    for extra_topic in extra_topics:
        file_name = os.path.join(file_path, extra_topic, name + ".jpg")
        img = cv2.imread(file_name)
        gbuff[extra_topic] = img[:, :, ::-1]
    gbuffers_list = [gbuff['SceneColor'], gbuff['SceneDepth'][:, :, 0][:, :, np.newaxis], gbuff['GBufferA'],
                     gbuff['GBufferB'], gbuff['GBufferC'], gbuff['GBufferD'],
                     gbuff['GBufferSSAO'][:, :, 0][:, :, np.newaxis], gbuff['CustomStencil'][:, :, 0][:, :, np.newaxis]]
    stacked_image = np.concatenate(gbuffers_list, axis=2)
    return stacked_image


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--data-path', default="/home/adt/bags/work_space/datasets",
                        help='your data root for carla')
    args = parser.parse_args()
    main(args)
